package com.iwuyc.tools.commons.concurrency.impl;

import com.iwuyc.tools.commons.concurrency.ConcurrencyTokenService;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.util.Optional;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
public class ConcurrencyTokenServiceImpl implements ConcurrencyTokenService {
    private final AtomicLong tokenVersionManager = new AtomicLong();
    private final BlockingQueue<Long> tokenBucket;
    private final int maxConcurrency;
    @Getter
    private int concurrency;

    /**
     * 创建一个令牌桶的管理类
     *
     * @param maxConcurrency 最大的并发数
     */
    public ConcurrencyTokenServiceImpl(int maxConcurrency) {
        this(maxConcurrency, 50);
    }

    public ConcurrencyTokenServiceImpl(int maxConcurrency, int concurrency) {
        this.maxConcurrency = maxConcurrency;
        this.tokenBucket = new ArrayBlockingQueue<>(maxConcurrency);

        incrementAndFillToken(concurrency);
    }

    private void incrementAndFillToken(int concurrency) {
        if (this.maxConcurrency < concurrency) {
            throw new IllegalArgumentException("maxConcurrency不允许比concurrency小");
        }
        this.tokenBucket.clear();
        this.concurrency = concurrency;
        long tokenVersion = this.tokenVersionManager.incrementAndGet();
        for (int i = 0; i < this.concurrency; i++) {
            this.tokenBucket.add(tokenVersion);
        }
    }

    @Override
    public Optional<ConcurrencyToken> take(long timeout, TimeUnit timeoutUnit) throws InterruptedException {

        final Long tokenVersion = this.tokenBucket.poll(timeout, timeoutUnit);
        if (null == tokenVersion) {
            return Optional.empty();
        }
        final ConcurrencyTokenImpl concurrencyToken = new ConcurrencyTokenImpl(this, tokenVersion);
        return Optional.of(concurrencyToken);
    }

    void feedbackToken(long tokenVersion) {
        final long currencyVersion = this.tokenVersionManager.get();
        if (tokenVersion == currencyVersion && !this.tokenBucket.offer(tokenVersion)) {
            log.info("归还token失败。");
        } else {
            log.info("归还token成功。");
        }
    }

    @Override
    public boolean concurrency(int concurrency) {
        this.incrementAndFillToken(concurrency);
        return true;
    }


    public static class ConcurrencyTokenImpl implements ConcurrencyTokenService.ConcurrencyToken {
        private final ConcurrencyTokenServiceImpl concurrencyTokenService;
        private final long tokenVersion;

        public ConcurrencyTokenImpl(ConcurrencyTokenServiceImpl concurrencyTokenService, long tokenVersion) {
            this.concurrencyTokenService = concurrencyTokenService;
            this.tokenVersion = tokenVersion;
        }

        @Override
        public void close() {
            this.concurrencyTokenService.feedbackToken(this.tokenVersion);
        }
    }
}
